package datastructure.avltree;


import java.util.*;

/**
 * 平衡搜索树--让二分搜索树保持平衡
 * 平衡因子：左右子树的高度差的绝对值，，一旦有个结点的平衡因子大于1，则不满足平衡二叉树
 *
 * @param <K> 泛型
 */
public class AVLTree<K extends Comparable<K>, V> {
    private int size;
    private Node root;

    private class Node {
        K key;
        V value;
        int height;
        Node left;
        Node right;

        public Node(K key, V value) {
            this.key = key;
            this.value = value;
            left = null;
            right = null;
            height = 1;
        }
    }

    private int getHeight(Node node) {
        if (node == null) return 0;
        return node.height;
    }

    public AVLTree() {
        size = 0;
        root = null;
    }

    /**
     * 添加元素
     *
     * @param k 元素
     */
    public void add(K k, V v) {
        root = add(root, k, v);
    }

    private Node add(Node node, K k, V v) {
        if (node == null) {
            size++;
            return new Node(k, v);
        }
        if (k.compareTo(node.key) < 0) {
            node.left = add(node.left, k, v);
        } else if (k.compareTo(node.key) > 0) {
            node.right = add(node.right, k, v);
        }

        // 更新height
        node.height = 1 + Math.max(getHeight(node.left), getHeight(node.right));

        // 计算平衡因子
        int balanceFactor = getBalanceFactor(node);
        if (Math.abs(balanceFactor) > 1) {
            System.out.println("该结点不平衡 " + balanceFactor);
        }

        // 平衡因子维护
        if (balanceFactor > 1 && getBalanceFactor(node.left) >= 0) {
            //执行右旋转
            return rightRotate(node);
        }
        if (balanceFactor < -1 && getBalanceFactor(node.right) <= 0) {
            return leftRotate(node);
        }
        //LR
        if (balanceFactor > 1 && getBalanceFactor(node.left) < 0) {
            node.left = leftRotate(node.left);
            return rightRotate(node);
        }
        //RL
        if (balanceFactor < -1 && getBalanceFactor(node.left) > 0) {
            node.right = leftRotate(node.right);
            return rightRotate(node);
        }

        return node;
    }

    private Node leftRotate(Node y) {
        Node x = y.left;
        Node t2 = x.right;

        //向右旋转的过程：
        x.left = y;
        y.right = t2;
        // 更新xy的高度
        y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
        x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
        return x;
    }

    private Node rightRotate(Node y) {
        Node x = y.left;
        Node t3 = x.right;

        //向右旋转的过程：
        x.right = y;
        y.right = t3;
        // 更新xy的高度
        y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
        x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
        return x;
    }

    // 判断是不是二分搜索树
    public boolean isBST() {
        List<K> keys = new ArrayList<>();
        inOrder(root, keys);
        for (int i = 1; i < keys.size(); i++) {
            if (keys.get(i - 1).compareTo(keys.get(i)) > 0) {
                return false;
            }
        }
        return true;
    }

    // 中序遍历把元素都拿出来
    private void inOrder(Node node, List<K> keys) {
        if (node == null) {
            return;
        }
        inOrder(node.left, keys);
        keys.add(node.key);
        inOrder(node.right, keys);

    }

    // 判断一颗二叉树是否是平衡二叉树
    public boolean isBalanced() {
        return isBalanced(root);
    }

    private boolean isBalanced(Node node) {
        // 左右子树的高度值不超过1
        if (node == null) {
            return true;
        }
        int balanceFactor = getBalanceFactor(node);
        if (Math.abs(balanceFactor) > 1) {
            return false;
        } else {
            return isBalanced(node.left) && isBalanced(node.right);
        }

    }

    private int getBalanceFactor(Node node) {
        if (node == null) {
            return 0;
        }
        return Math.abs(getHeight(node.left) - getHeight(node.right));
    }

    /**
     * 是否包含元素e
     *
     * @param k 元素e
     * @return 是否有
     */
    public boolean contains(K k) {
        return contains(root, k);
    }

    private boolean contains(Node root, K k) {
        if (root == null) {
            return false;
        }
        if (k.compareTo(root.key) < 0) {
            return contains(root.left, k);
        } else if (k.equals(root.key)) {
            return true;
        } else {
            return contains(root.left, k);
        }
    }

    /**
     * 前序遍历
     */
    public void preOrder() {
        preOrder(root);
    }

    private void preOrder(Node root) {
        if (root == null) {
            return;
        }
        System.out.println(root.key);
        preOrder(root.left);
        preOrder(root.right);
    }

    /**
     * 后序遍历
     */
    public void postOrder() {
        postOrder(root);
    }

    private void postOrder(Node root) {
        if (root == null) {
            return;
        }
        postOrder(root.left);
        postOrder(root.right);
        System.out.println(root.key);
    }
//    private void add(Node root, E e) {
//        if (e.equals(root.e)){
//            return;
//        }else if (e.compareTo(root.e) < 0 && root.left == null){
//            root.left = new Node(e);
//        } else if (e.compareTo(root.e) > 0 && root.right == null){
//            root.right = new Node(e);
//        }
//        if (e.compareTo(root.e) < 0){
//            add(root.left, e);
//        } else {
//            add(root.right, e);
//        }
//    }

    /**
     * 二分搜索树的前序遍历，非递归实现
     */
    public void preOrderNR() {
        Stack<Node> stack = new Stack<>();
        if (root != null) {
            stack.push(root);
        }
        preOrderNR(stack);
    }

    private void preOrderNR(Stack<Node> stack) {
        // 栈为空时作为终止条件
        while (!stack.isEmpty()) {
            Node node = stack.pop();
            // 每次出栈，先压入非空的右左孩子（注意顺序）
            System.out.println(node.key);
            if (node.right != null) {
                stack.push(node.right);
            }
            if (node.left != null) {
                stack.push(node.left);
            }
        }
    }

    /**
     * 二叉树的层序遍历
     */
    public void orderInLine() {
        // 先定义个队列
        Queue<Node> queue = new LinkedList<>();
        queue.add(root);
        while (!queue.isEmpty()) {
            Node cur = queue.remove();
            // 移除队首元素后，把改结点的左右孩子入队尾，注意顺序啊，先左后右
            System.out.println(cur.key);
            if (cur.left != null)
                queue.add(cur.left);
            if (cur.right != null)
                queue.add(cur.right);
        }
    }

    /**
     * 寻找最小元素
     *
     * @return 最小元素
     */
    public K minElement() {
        if (root == null) {
            throw new IllegalArgumentException("二分搜索树为空");
        }
        Node min = minElement(root);
        return min.key;
    }

    private Node minElement(Node root) {
        if (root.left == null) {
            return root;
        }
        return minElement(root.left);
    }

    /**
     * 寻找最大元素
     *
     * @return 最大元素
     */
    public K maxElement() {
        if (root == null) {
            throw new IllegalArgumentException("二分搜索树为空");
        }
        Node max = maxElement(root);
        return max.key;
    }

    private Node maxElement(Node root) {
        if (root.right == null) {
            return root;
        }
        return minElement(root.right);
    }

    public K removeMin() {
        K k = minElement();
        root = removeMin(root);
        return k;
    }

    private Node removeMin(Node node) {
        if (node.left == null) {
            K k = node.key;
            Node right = node.right;
            node.right = null;
            size--;
            return right;
        }
        K k = node.key;
        node.left = removeMin(node.left);
        return node;
    }

    public K removeMax() {
        K k = maxElement();
        root = removeMax(root);
        return k;
    }

    private Node removeMax(Node node) {
        if (node.right == null) {
            Node left = node.left;
            node.left = null;
            size--;
            return left;
        }
        node.right = removeMax(node.right);
        return node;
    }

    public V removeElem(K k) {
        if (!contains(k)) {
            return null;
        }
        return removeElem(root, k).value;
    }

    private Node removeElem(Node node, K k) {
        if (node == null) {
            return null;
        }
        if (k.compareTo(node.key) < 0) {
            node.left = removeElem(node.left, k);
            return node;
        } else if (k.compareTo(node.key) > 0) {
            node.right = removeElem(node.right, k);
            return node;
        } else {
            if (node.left == null) {
                Node right = node.right;
                node.right = null;
                size--;
                return right;
            }
            if (node.right == null) {
                Node left = node.left;
                node.left = null;
                size--;
                return left;
            }
            // 左右都不为空
            Node successor = minElement(node.right);
            successor.right = removeMin(node.right);
            successor.left = node.left;

            node.left = node.right = null;
            return successor;
        }
    }

    /**
     * 是否为空
     *
     * @return 是否为空
     */
    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * 获取元素个数
     *
     * @return 元素个数
     */
    public int getSize() {
        return size;
    }


}

